const fs = require('fs').promises;
const path = require('path');
const xml2js = require('xml2js');
const BaseDatasetWrapper = require('../../BaseDatasetWrapper');
const VOCLabelLoader = require('./VOCLabelLoader');
const { DatasetListDB } = require('../../../dbs/dbDatasetList');
const { getAllFilesInDirectory } = require('../../../utils/fileUtils');
const sharp = require('sharp');

class VOCDatasetWrapper extends BaseDatasetWrapper {
    // 读入目录信息，解析对应的标注结果
    async parse(dataset_info) {
        // console.log('VOCDatasetWrapper::parse dataset_info = ', dataset_info);
        const datasetId = dataset_info.dataset_id;
        const datasetDir = dataset_info.dataset_dir;
        this.parseCallback(datasetId, 'parse_status', '读取图片列表...');
        let imageFiles = await getAllFilesInDirectory(datasetDir, ['.jpg', '.jpeg', '.png', '.bmp'], ['workspace']);
        this.parseCallback(datasetId, 'parse_status', '读取标注文件列表...');
        let labelFilePaths = await getAllFilesInDirectory(datasetDir, ['.xml'], ['workspace']);
        this.parseCallback(datasetId, 'parse_start', imageFiles.length);

        // 创建映射关系，获取categories信息
        let labelMap = {};
        let categories_name = new Set();

        await Promise.all(
            labelFilePaths.map(async (labelPath) => {
                const labelFileName = path.basename(labelPath);
                labelMap[labelFileName] = labelPath;
                // console.log('labelPath = ', labelPath);
                let categories = await this.getCategories(labelPath);
                categories.forEach((name) => categories_name.add(name));
            })
        );

        // 创建新的目录用于存放复制的标注文件
        const workDir = path.join(datasetDir, 'workspace', datasetId);
        await fs.mkdir(workDir, { recursive: true });

        // 加载每一张图片对应的标注数据
        let totalSamples = imageFiles.length;
        let labeledSamples = 0;
        let matchingPairs = [];
        for (let i = 0; i < imageFiles.length; i++) {
            let imageFile = imageFiles[i];

            // 获取标注文件的路径
            const labelFileName = path.basename(imageFile, path.extname(imageFile)) + '.xml';
            const labelFilePath = labelMap[labelFileName];

            // 获取图片相对于 dataset_dir 的相对路径
            const imageRelativePath = path.relative(datasetDir, imageFile);
            const labelRelativeDir = path.dirname(imageRelativePath);
            const newLabelFileDir = path.join(workDir, labelRelativeDir);
            const newLabelFilePath = path.join(newLabelFileDir, labelFileName);

            // 确保目标目录存在，并复制标注文件到新的目录中
            await fs.mkdir(newLabelFileDir, { recursive: true });

            // 拷贝或者新建一个标注文件
            if (labelFilePath) {
                await fs.copyFile(labelFilePath, newLabelFilePath);
                labeledSamples++;
            } else {
                await fs.writeFile(newLabelFilePath, '<annotation></annotation>', 'utf8');
            }

            // 保存图像和标注的配对文件
            matchingPairs.push({
                image_path: path.relative(datasetDir, imageFile),
                label_path: path.relative(datasetDir, newLabelFilePath)
            });

            // 更新数据加载进展
            this.parseCallback(datasetId, 'parse_percent', (((i + 1) * 100) / totalSamples).toFixed(2));
        }

        // 保存matching pairs结果
        await this.saveMatchingPairs(datasetId, matchingPairs);

        // 汇总图片目录名称，方便后续切分训练集和验证集使用
        let imageDirs = new Set();
        imageFiles.forEach((file) => imageDirs.add(path.dirname(file)));
        await this.createSubsets(dataset_info, Array.from(imageDirs));

        // 保存categories信息
        // console.log(categories_name);
        let categories = Array.from(categories_name).map((category, index) => {
            return {
                id: index,
                name: category,
                supercategory: category
            };
        });
        this.saveCategories(dataset_info.dataset_id, categories);

        return {
            code: 0,
            msg: {
                totalSamples: totalSamples,
                labeledSamples: labeledSamples
            }
        };
    }

    // 从voc的xml文件中获取category信息，有可能导致信息比较慢
    async getCategories(label_path) {
        try {
            const xmlData = await fs.readFile(label_path, 'utf8');
            const xmlRoot = await xml2js.parseStringPromise(xmlData);
            const annotation = xmlRoot.annotation;
            let categories = {};
            for (let i = 0; i < annotation.object.length; i++) {
                let object = annotation.object[i];
                if (object.name) {
                    categories[object.name] = 1;
                }
            }
            return Object.keys(categories);
        } catch (error) {
            console.error(`Error load voc label file ${label_path},  error: `, error);
            return [];
        }
    }

    // 导出数据集标注结果列表
    async export(toFormat, dataset_info, subsets_info) {
        // console.log('VOCDatasetWrapper::export toFormat = ', toFormat);
        // console.log('VOCDatasetWrapper::export dataset_info = ', dataset_info);
        // console.log('VOCDatasetWrapper::export subsets_info = ', subsets_info);

        // 将train、val、test的子集目录，以及子集的图片文件列表，写入到本地文件中
        let dataset_dir = dataset_info.dataset_dir;
        let workDir = path.join(dataset_dir, 'workspace', dataset_info.id);
        for (let i = 0; i < subsets_info.length; i++) {
            let subset = subsets_info[i];
            let imageDir = path.join(dataset_dir, subset.image_dir);
            let imageFiles = await getAllFilesInDirectory(imageDir, ['.jpg', '.jpeg', '.png', '.bmp'], ['workspace']);
            let output_file = path.join(workDir, `${subset.subset_type}.txt`);

            // 判断是否要转换图片格式到jpg
            let needConversion = imageFiles.some((file) => !['.jpg', '.jpeg'].includes(path.extname(file).toLowerCase()));
            if (needConversion) {
                let convertedImagesDir = path.join(workDir, 'converted_images');
                await fs.mkdir(convertedImagesDir, { recursive: true });

                for (let file of imageFiles) {
                    let ext = path.extname(file).toLowerCase();
                    let fileNameWithoutExt = path.parse(file).name;
                    let destPath = path.join(convertedImagesDir, `${fileNameWithoutExt}.jpg`);
                    // Check if the destination file already exists
                    if (await fs.exists(destPath)) {
                        // console.log(`File ${destPath} already exists, skipping conversion.`);
                        continue;
                    }

                    if (ext !== '.jpg' && ext !== '.jpeg') {
                        // Convert to JPEG
                        await sharp(file).jpeg({ quality: 90 }).toFile(destPath);
                        console.log(`Converted ${file} to ${destPath}`);
                    } else {
                        // Copy JPEG without conversion
                        await fs.copyFile(file, destPath);
                        console.log(`Copied ${file} to ${destPath}`);
                    }
                }

                // Update image directory to point to converted images
                subset.image_dir = path.relative(dataset_dir, convertedImagesDir);
            }

            // Write the basenames of the image files to the output file
            const fileContents = imageFiles.map((file) => path.parse(file).name).join('\n');
            await fs.writeFile(output_file, fileContents);
            console.log(`Written ${imageFiles.length} filenames to ${output_file}`);

            // store output file path to subset info
            subset.label_list_path = path.relative(dataset_dir, output_file);
        }
    }

    // 转换数据集到指定格式
    async convert(toFormat, dataset_info) {
        let format = toFormat.toLowerCase();
        if (format === 'voc') {
            return true;
        }
        const datasetId = dataset_info.id;
        const datasetDir = dataset_info.dataset_dir;

        try {
            // 获取数据集的标注列表
            let dblist = this.app.list_dbs[datasetId];
            if (!dblist) {
                this.app.list_dbs[datasetId] = new DatasetListDB(this.app, datasetId);
                dblist = this.app.list_dbs[datasetId];
            }
            let pairs = await dblist.pairs.getPairs();

            // 动态加载labelloader类
            const srcLoader = new VOCLabelLoader(this.app);
            const loaderFile = toFormat.toUpperCase() + 'LabelLoader.js';
            const loaderClass = require(path.join(__dirname, loaderFile));
            const dstloader = new loaderClass(this.app);

            // 对于每一个pair进行转换
            let categories = await this.app.db.categories.getCategories(datasetId);
            for (let i = 0; i < pairs.length; i++) {
                let pair = pairs[i];
                pair.image_path = path.join(datasetDir, pair.image_path);
                pair.label_path = path.join(datasetDir, pair.label_path);
                let label_data = await srcLoader.load({
                    dataset_id: datasetId,
                    image_path: pair.image_path,
                    label_path: pair.label_path
                });
                let label_wrapper = this.wrapLabelData(dataset_info, pair, label_data);
                let label_info = await dstloader.toString(label_wrapper, categories);

                // 把结果保存到本地文件中
                let outputPath = BaseDatasetWrapper.getConvertedPath(toFormat, dataset_info, pair);
                await fs.writeFile(outputPath, label_info);
            }
            return { code: 0, msg: `VOC to ${toFormat.toUpperCase()} successfully` };
        } catch (e) {
            console.error('VOCDatasetWrapper::convert error = ', e);
            return { code: -1, msg: e.message };
        }
    }
}

module.exports = VOCDatasetWrapper;
